import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from tqdm import tqdm
import pickle

# 1. 加载手写数字数据集
digits = load_digits()
X = digits.data
y = digits.target

# 2. 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)  # 设置random_state确保可复现

# 3. 尝试不同的K值（1~40），记录准确率
k_values = range(1, 41)
accuracies = []
for k in tqdm(k_values, desc="Finding optimal K"):
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(X_train, y_train)
    y_pred = knn.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    accuracies.append(accuracy)

# 4. 找出最优的K值和对应的准确率
optimal_k = k_values[np.argmax(accuracies)]
optimal_accuracy = max(accuracies)
print(f"最优的K值是 {optimal_k}，对应的准确率是 {optimal_accuracy:.4f}")

# 5. 训练最优K值的KNN模型并保存
best_knn = KNeighborsClassifier(n_neighbors=optimal_k)
best_knn.fit(X_train, y_train)
with open("best_knn_model.pkl", "wb") as f:
    pickle.dump(best_knn, f)

# 6. 绘制准确率变化折线图并保存为PDF
plt.figure(figsize=(10, 6))
plt.plot(k_values, accuracies, marker="o", linestyle="-")
plt.axvline(x=optimal_k, color="red", linestyle="--")
plt.text(optimal_k + 0.5, optimal_accuracy, f"k={optimal_k}, Accuracy={optimal_accuracy:.4f}", color="red")
plt.xlabel("K Value")
plt.ylabel("Accuracy")
plt.title("Accuracy of different k values")
plt.grid(True)
plt.savefig("accuracy_plot.pdf")
plt.close()